#!/usr/bin/python
# -*-coding:utf-8-*-
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
# plt.switch_backend('agg')
import seaborn as sns

plt.rcParams[u'font.sans-serif'] = ['simhei']
plt.rcParams['axes.unicode_minus'] = False

import os
import pandas as pd
import numpy as np
from scipy import stats

# TODO - 读取数据
def get_group_analysis_plot(factor_name,
                            group_ret_data_dir,
                            group_indicator_data_dir,
                            analysis_result_output_data_dir,
                            show_top_bottom_only=True):
    # ic
    ic_data = pd.read_excel(os.path.join(group_ret_data_dir, factor_name+'_ic_data.xlsx'))
    ic_data = ic_data.set_index('date')

    ic_data = ic_data.resample('1M')

    used_ic_data = ic_data[['ic']].mean() / ic_data[['ic']].std()

    used_ic_data = used_ic_data.rename(columns = {'ic': 'icir'})

    used_ic_data['ma6'] = used_ic_data['icir'].rolling(6).mean()

    # long-short
    long_short_data = pd.read_excel(os.path.join(group_ret_data_dir, factor_name+'_group_abs_ret_data.xlsx'))

    long_short_data['ls'] = long_short_data.iloc[:, -1] - long_short_data.iloc[:, 1]

    long_short_data = long_short_data.set_index('date')

    long_short_data = long_short_data.resample('1M')

    used_long_short_data = long_short_data[['ls']].mean() * 100

    used_long_short_data['ma6'] = used_long_short_data['ls'].rolling(6).mean()

    # used_long_short_data['label'] = 0.0
    used_long_short_data.index = used_long_short_data.index.strftime('%Y-%m')

    def group_analysis_result_process(df):
        df = df.set_index(['selection_date', 'group'])

        df = df.to_panel()

        df = df.swapaxes(axis1=0, axis2=2)

        df = df.to_frame(filter_observations=False)

        df = df.reset_index('selection_date').set_index('selection_date')

        df.columns = ['g%s' % g for g in df.columns]

        return df

    # volatility analysis
    volatility_analysis = pd.read_excel(os.path.join(group_indicator_data_dir, factor_name+'_group_volatility_analysis.xlsx'))
    volatility_analysis = group_analysis_result_process(volatility_analysis)
    volatility_analysis['spread'] = volatility_analysis.iloc[:, -1] - volatility_analysis.iloc[:, 0]

    # turnover analysis
    turnover_analysis = pd.read_excel(os.path.join(group_indicator_data_dir, factor_name+'_group_turnover_analysis.xlsx'))
    turnover_analysis = group_analysis_result_process(turnover_analysis)
    turnover_analysis['spread'] = turnover_analysis.iloc[:, -1] - turnover_analysis.iloc[:, 0]

    # profitability analysis
    profitability_analysis = pd.read_excel(os.path.join(group_indicator_data_dir, factor_name+'_group_profitability_analysis.xlsx'))
    profitability_analysis = group_analysis_result_process(profitability_analysis)
    profitability_analysis['spread'] = profitability_analysis.iloc[:, -1] - profitability_analysis.iloc[:, 0]

    # mom analysis
    mom_analysis = pd.read_excel(os.path.join(group_indicator_data_dir, factor_name+'_group_mom_analysis.xlsx'))
    mom_analysis = group_analysis_result_process(mom_analysis)
    mom_analysis['spread'] = mom_analysis.iloc[:, -1] - mom_analysis.iloc[:, 0]

    # growth analysis
    growth_analysis = pd.read_excel(os.path.join(group_indicator_data_dir, factor_name+'_group_growth_analysis.xlsx'))
    growth_analysis = group_analysis_result_process(growth_analysis)
    growth_analysis['spread'] = growth_analysis.iloc[:, -1] - growth_analysis.iloc[:, 0]

    # beta analysis
    beta_analysis = pd.read_excel(os.path.join(group_indicator_data_dir, factor_name+'_group_beta_analysis.xlsx'))
    beta_analysis = group_analysis_result_process(beta_analysis)
    beta_analysis['spread'] = beta_analysis.iloc[:, -1] - beta_analysis.iloc[:, 0]

    # valuation analysis
    valuation_analysis = pd.read_excel(os.path.join(group_indicator_data_dir, factor_name+'_group_valuation_analysis.xlsx'))
    valuation_analysis = group_analysis_result_process(valuation_analysis)
    valuation_analysis['spread'] = valuation_analysis.iloc[:, -1] - valuation_analysis.iloc[:, 0]

    # TODO - 画热力图
    ## 超额收益热力图
    fig = plt.figure(figsize=(40, 15))  # 设置画面大小
    ax1 = fig.add_subplot(421)
    ax2 = fig.add_subplot(422)
    ax3 = fig.add_subplot(423)
    ax4 = fig.add_subplot(424)
    ax5 = fig.add_subplot(425)
    ax6 = fig.add_subplot(426)
    ax7 = fig.add_subplot(427)
    ax8 = fig.add_subplot(428)

    # used_ic_data.plot(ax=ax1)
    used_long_short_data.plot(ax=ax1, kind='bar')

    # ax1.set_title('%s - monthly ICIR' % (factor_name))
    ax1.set_title('%s - monthly long-short' % (factor_name))

    if show_top_bottom_only:
        valuation_analysis = valuation_analysis.iloc[:, [0, -2, -1]]

    valuation_analysis.iloc[:, :-1].plot(ax=ax2)
    ax22 = ax2.twinx()
    valuation_analysis.iloc[:, -1].plot(ax=ax22, color='red')

    ax2.set_title('%s - valuation analysis' % (factor_name))

    if show_top_bottom_only:
        # profitability_analysis = profitability_analysis.iloc[:, [0, -1]]
        profitability_analysis = profitability_analysis.iloc[:, [0, -2, -1]]

    # profitability_analysis.plot(ax=ax3)
    profitability_analysis.iloc[:, :-1].plot(ax=ax3)
    ax32 = ax3.twinx()
    profitability_analysis.iloc[:, -1].plot(ax=ax32, color='red')

    ax3.set_title('%s - profitability analysis' % (factor_name))

    if show_top_bottom_only:
        # growth_analysis = growth_analysis.iloc[:, [0, -1]]
        growth_analysis = growth_analysis.iloc[:, [0, -2, -1]]

    # growth_analysis.plot(ax=ax4)
    growth_analysis.iloc[:, :-1].plot(ax=ax4)
    ax42 = ax4.twinx()
    growth_analysis.iloc[:, -1].plot(ax=ax42, color='red')

    ax4.set_title('%s - growth analysis' % (factor_name))

    if show_top_bottom_only:
        # mom_analysis = mom_analysis.iloc[:, [0, -1]]
        mom_analysis = mom_analysis.iloc[:, [0, -2, -1]]

    # mom_analysis.plot(ax=ax5)
    mom_analysis.iloc[:, :-1].plot(ax=ax5)
    ax52 = ax5.twinx()
    mom_analysis.iloc[:, -1].plot(ax=ax52, color='red')

    ax5.set_title('%s - mom analysis' % (factor_name))

    if show_top_bottom_only:
        # beta_analysis = beta_analysis.iloc[:, [0, -1]]
        beta_analysis = beta_analysis.iloc[:, [0, -2, -1]]

    # beta_analysis.plot(ax=ax6)
    beta_analysis.iloc[:, :-1].plot(ax=ax6)
    ax62 = ax6.twinx()
    beta_analysis.iloc[:, -1].plot(ax=ax62, color='red')

    ax6.set_title('%s - beta analysis' % (factor_name))


    if show_top_bottom_only:
        # volatility_analysis = volatility_analysis.iloc[:, [0, -1]]
        volatility_analysis = volatility_analysis.iloc[:, [0, -2, -1]]

    # volatility_analysis.plot(ax=ax7)
    volatility_analysis.iloc[:, :-1].plot(ax=ax7)
    ax72 = ax7.twinx()
    volatility_analysis.iloc[:, -1].plot(ax=ax72, color='red')

    ax7.set_title('%s - volatility analysis' % (factor_name))

    if show_top_bottom_only:
        # turnover_analysis = turnover_analysis.iloc[:, [0, -1]]
        turnover_analysis = turnover_analysis.iloc[:, [0, -2, -1]]

    # turnover_analysis.plot(ax=ax8)
    turnover_analysis.iloc[:, :-1].plot(ax=ax8)
    ax82 = ax8.twinx()
    turnover_analysis.iloc[:, -1].plot(ax=ax82, color='red')

    ax8.set_title('%s - turnover analysis' % (factor_name))

    plt.tight_layout()

    plt.savefig(os.path.join(analysis_result_output_data_dir, factor_name + '_group_analysis_stats.png'))

    plt.close('all')

    print(factor_name, 'done!')


def get_long_short_ic_heatmap(factor_name,
                               group_ret_data_dir,
                               analysis_result_output_data_dir):
    # ic
    ic_data = pd.read_excel(os.path.join(group_ret_data_dir, factor_name+'_ic_data.xlsx'))
    ic_data = ic_data.set_index('date')

    ic_data = ic_data.resample('1M')

    used_ic_data = ic_data[['ic']].mean()

    used_ic_data['year'] = used_ic_data.index.year
    used_ic_data['month'] = used_ic_data.index.month

    used_ic_data = used_ic_data.groupby(['year', 'month'])[['ic']].sum()

    used_ic_data = used_ic_data.to_panel()
    used_ic_data = used_ic_data.swapaxes(axis1=0, axis2=2)

    used_ic_data = used_ic_data.to_frame(filter_observations=False)
    used_ic_data.index = used_ic_data.index.droplevel('minor')

    used_ic_data = (used_ic_data * 100).round(3)

    used_icir_data = ic_data[['ic']].mean() / ic_data[['ic']].std()

    used_icir_data = used_icir_data.rename(columns = {'ic': 'icir'})

    used_icir_data['year'] = used_icir_data.index.year
    used_icir_data['month'] = used_icir_data.index.month

    used_icir_data = used_icir_data.groupby(['year', 'month'])[['icir']].sum()

    used_icir_data = used_icir_data.to_panel()
    used_icir_data = used_icir_data.swapaxes(axis1=0, axis2=2)

    used_icir_data = used_icir_data.to_frame(filter_observations=False)
    used_icir_data.index = used_icir_data.index.droplevel('minor')

    used_icir_data = used_icir_data.round(3)

    # long-short
    long_short_data = pd.read_excel(os.path.join(group_ret_data_dir, factor_name+'_group_abs_ret_data.xlsx'))

    long_short_data = long_short_data.set_index('date')

    long_short_data['ls'] = long_short_data.iloc[:, -1] - long_short_data.iloc[:, 1]

    long_short_data = long_short_data.resample('1M')

    used_long_short_data = long_short_data[['ls']].sum()

    used_long_short_data['year'] = used_long_short_data.index.year
    used_long_short_data['month'] = used_long_short_data.index.month

    used_long_short_data = used_long_short_data.groupby(['year', 'month'])[['ls']].sum()

    used_long_short_data = used_long_short_data.to_panel()
    used_long_short_data = used_long_short_data.swapaxes(axis1=0, axis2=2)

    used_long_short_data = used_long_short_data.to_frame(filter_observations=False)
    used_long_short_data.index = used_long_short_data.index.droplevel('minor')

    # used_long_short_data = used_long_short_data.round(3)
    used_long_short_data = (used_long_short_data * 100).round(3)

    # TODO - 画热力图
    ## 超额收益热力图
    fig = plt.figure(figsize=(20, 10))  # 设置画面大小
    ax1 = fig.add_subplot(131)
    ax2 = fig.add_subplot(132)
    ax3 = fig.add_subplot(133)

    sns.heatmap(used_ic_data,
                annot=True,
                vmax=1,
                fmt='.3',
                center=0.0,
                annot_kws={'size': 10, 'weight': 'bold'},
                square=True,
                ax=ax1,
                cmap="Reds")

    ax1.set_title('IC(%)')

    sns.heatmap(used_icir_data,
                annot=True,
                vmax=1,
                fmt='.3',
                center=0.0,
                annot_kws={'size': 10, 'weight': 'bold'},
                square=True,
                ax=ax2,
                cmap="Reds")

    ax2.set_title('ICIR(%)')

    sns.heatmap(used_long_short_data,
                annot=True,
                vmax=1,
                fmt='.3',
                center=0.0,
                annot_kws={'size': 10, 'weight': 'bold'},
                square=True,
                ax=ax3,
                cmap="Reds")

    ax3.set_title('Long-Short(%)')

    plt.tight_layout()

    plt.savefig(os.path.join(analysis_result_output_data_dir, factor_name + '_icir_ls_stats.png'))

    plt.close('all')
    print(factor_name, 'done!')

    return_used_long_short_data = used_long_short_data.copy()
    return_used_icir_data = used_icir_data.copy()
    return_used_ic_data = used_ic_data.copy()

    return_used_long_short_data['factor_name'] = factor_name
    return_used_icir_data['factor_name'] = factor_name
    return_used_ic_data['factor_name'] = factor_name

    return return_used_long_short_data, return_used_icir_data, return_used_ic_data


if __name__ == '__main__':
    # group_ret_data_dir = '/db/zg_data/zbc/factor_analysis/monitor/ew_v1/20200305/ths_ew_profit_06_rank'
    # group_indicator_data_dir = '/db/zg_data/zbc/factor_analysis/monitor/ew_v1/20200305/ths_ew_profit_06_rank/indicator'

    # raw_data_dir = '/db/zg_data/zbc/factor_analysis/monitor/ew_v1/20200305'
    raw_data_dir = '/db/zg_data/zbc/factor_analysis/monitor/ew_real/20200316'

    analysis_result_output_data_dir = raw_data_dir + '/analysis_plot'
    icir_analysis_result_output_data_dir = raw_data_dir + '/analysis_plot/icir'

    if not os.path.exists(icir_analysis_result_output_data_dir):
        os.makedirs(icir_analysis_result_output_data_dir)

    factor_name_list = [fn for fn in os.listdir(raw_data_dir) if 'analysis_plot' != fn]

    long_short_data_summary = []
    ic_data_summary = []
    icir_data_summary = []
    for factor_name in factor_name_list:
        group_ret_data_dir =  raw_data_dir + '/%s' % (factor_name)
        group_indicator_data_dir = raw_data_dir + '/%s/indicator' % (factor_name)

        try:
            get_group_analysis_plot(factor_name,
                                    group_ret_data_dir,
                                    group_indicator_data_dir,
                                    analysis_result_output_data_dir,
                                    show_top_bottom_only=True)

            return_used_long_short_data, return_used_icir_data, return_used_ic_data = \
                get_long_short_ic_heatmap(factor_name,
                                          group_ret_data_dir,
                                          icir_analysis_result_output_data_dir)

            long_short_data_summary.append(return_used_long_short_data)
            icir_data_summary.append(return_used_icir_data)
            ic_data_summary.append(return_used_ic_data)
        except:
            print(factor_name, 'exception!')

    long_short_data_summary = pd.concat(long_short_data_summary, axis=0)
    icir_data_summary = pd.concat(icir_data_summary, axis=0)
    ic_data_summary = pd.concat(ic_data_summary, axis=0)

    long_short_data_summary.to_excel(os.path.join(icir_analysis_result_output_data_dir, 'long_short_data_summary.xlsx'))
    icir_data_summary.to_excel(os.path.join(icir_analysis_result_output_data_dir, 'icir_data_summary.xlsx'))
    ic_data_summary.to_excel(os.path.join(icir_analysis_result_output_data_dir, 'ic_data_summary.xlsx'))

